#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os, re, argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

def hampel(x, k=5, n_sigmas=3.0):
    x = np.asarray(x, float)
    if len(x) < 2*k+1: return x
    med = pd.Series(x).rolling(2*k+1, center=True).median().to_numpy()
    diff = np.abs(x - med)
    mad  = pd.Series(diff).rolling(2*k+1, center=True).median().to_numpy()
    thr = n_sigmas * 1.4826 * mad
    mask = (diff > thr)
    y = x.copy()
    y[mask] = med[mask]
    return y

def ema(x, alpha=0.18):
    y = np.empty_like(x, dtype=float)
    y[0] = x[0]
    for i in range(1, len(x)):
        y[i] = (1 - alpha) * y[i-1] + alpha * x[i]
    return y

def detrend_ma(x, win=181):  # ≈2秒窗口(90Hz时)
    s = pd.Series(x, dtype=float)
    base = s.rolling(win, center=True, min_periods=1).mean()
    return (s - base).to_numpy()

def denoise_xyz(xyz, do_hampel=True, do_ema=True, alpha=0.18, do_detrend=False):
    X = xyz.copy().astype(float)
    for i in range(X.shape[1]):  # 对 X/Y/Z 各轴处理
        xi = X[:, i]
        if do_hampel: xi = hampel(xi, k=5, n_sigmas=3.0)         # 去尖峰
        if do_ema:    xi = ema(xi, alpha=alpha)                  # EMA 低通
        if do_detrend:xi = detrend_ma(xi, win=181)               # 可选：去慢漂移
        X[:, i] = xi
    return X

def slice_lines(df, start_line, end_line):
    if start_line is None and end_line is None:
        return df
    if start_line is None:
        start_line = 0
    if end_line is None or end_line > len(df):
        end_line = len(df)
    return df.iloc[start_line:end_line].reset_index(drop=True)


def read_motion_ly_or_csv(path):
    """
    兼容 Motion Analysis .ly（两行表头）与普通 CSV：
    - 自动定位 'Frame#' 行作为第一表头
    - 若下一行含 X1/Y1/Z1（或 X3/Y3/Z3），与上一行合并为完整列名
    - 之后的行按空白/逗号/分号分隔读取；不足的列用 NaN 补齐
    """
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        lines = [ln.rstrip("\n") for ln in f]

    # 找到含 'Frame#' 的行
    hdr1_idx = next((i for i, ln in enumerate(lines) if "Frame#" in ln), None)
    if hdr1_idx is None:
        # 普通 CSV 的情况
        for sep in [",", ";", r"\s+", "\t"]:
            try:
                df = pd.read_csv(path, sep=sep, engine="python")
                df = df.dropna(axis=1, how="all").reset_index(drop=True)
                if not df.empty:
                    return df
            except Exception:
                pass
        raise RuntimeError(f"无法读取文件：{path}")

    hdr1 = re.split(r"[,\t ]+", lines[hdr1_idx].strip())
    # 找到下一条非空行，看看是否是第二行子表头
    j = hdr1_idx + 1
    while j < len(lines) and not lines[j].strip():
        j += 1
    hdr2_tokens = re.split(r"[,\t ]+", lines[j].strip()) if j < len(lines) else []

    # 判断第二行是否为 X*/Y*/Z* 等
    has_xyz = any(re.fullmatch(r"[v|a]?[XYZ][0-9]+", tok) or re.fullmatch(r"[XYZ][0-9]+", tok) for tok in hdr2_tokens)
    if has_xyz:
        # 去掉第一行末尾的 Marker*，与第二行合并
        if hdr1[-1].lower().startswith("marker"):
            hdr = hdr1[:-1] + hdr2_tokens
        else:
            hdr = hdr1 + hdr2_tokens
        data_start = j + 1
    else:
        # 不是两行表头，就直接用第一行当列名
        hdr = hdr1
        data_start = hdr1_idx + 1

    # 解析数据行：按空白/逗号/分号切分，长度不够则补齐
    rows = []
    for ln in lines[data_start:]:
        if not ln.strip():
            continue
        toks = re.split(r"[,\t ]+", ln.strip())
        if len(toks) < len(hdr):
            toks = toks + [""] * (len(hdr) - len(toks))
        else:
            toks = toks[:len(hdr)]
        rows.append(toks)

    if not rows:
        raise RuntimeError(f"未在 {path} 中找到数据行（检查是否为空文件或权限）")

    df = pd.DataFrame(rows, columns=hdr)
    # 列转数值（不可转的保留 NaN）
    for c in df.columns:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    df = df.dropna(how="all", axis=1).reset_index(drop=True)
    return df

def find_time_column(df):
    """优先 t/time/timestamp/stamp 或 sec+{nsec,nanosec}；否则用样本索引。"""
    m = {str(c).lower(): c for c in df.columns}
    if "t" in m: return pd.to_numeric(df[m["t"]], errors="coerce").values, False
    for k in ["time", "timestamp", "stamp"]:
        if k in m: return pd.to_numeric(df[m[k]], errors="coerce").values, False
    if "sec" in m and ("nsec" in m or "nanosec" in m):
        sec = pd.to_numeric(df[m["sec"]], errors="coerce").fillna(0).values
        nk = "nsec" if "nsec" in m else "nanosec"
        nsec = pd.to_numeric(df[m[nk]], errors="coerce").fillna(0).values
        return (sec + nsec*1e-9), False
    return np.arange(len(df)), True

def pick_xyz(df, wanted):
    """大小写不敏感取列；若缺失则尽量猜：优先 X*/Y*/Z* 的任意编号，再兜底取前三列。"""
    m = {str(c).lower(): c for c in df.columns}
    cols = [m.get(name.lower()) for name in wanted]
    if all(c is not None for c in cols):
        return df[cols].apply(pd.to_numeric, errors="coerce").values

    # 尝试任意编号（比如 X1/Y1/Z1 或 X3/Y3/Z3）
    xs = [c for c in df.columns if re.fullmatch(r"[Xx]\d+", str(c))]
    ys = [c for c in df.columns if re.fullmatch(r"[Yy]\d+", str(c))]
    zs = [c for c in df.columns if re.fullmatch(r"[Zz]\d+", str(c))]
    if xs and ys and zs:
        return df[[xs[0], ys[0], zs[0]]].apply(pd.to_numeric, errors="coerce").values

    # 兜底：前三列
    if df.shape[1] >= 3:
        return df.iloc[:, :3].apply(pd.to_numeric, errors="coerce").values
    raise RuntimeError(f"未找到列 {wanted}，且列数不足 3。现有列名：{list(df.columns)}")
def trim(df, t, start, end):
    mask = np.ones(len(df), dtype=bool)
    if start is not None:
        mask &= (t >= start)
    if end is not None:
        mask &= (t <= end)
    return df[mask].reset_index(drop=True), t[mask]
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--start-line", type=int, default=None, help="起始行号（0 开始）")
    ap.add_argument("--end-line", type=int, default=None, help="结束行号（含）")
    ap.add_argument("--start", type=float, default=None, help="起始时间/帧号")
    ap.add_argument("--end", type=float, default=None, help="截止时间/帧号")
    ap.add_argument("face_file")
    ap.add_argument("arm_file")
    ap.add_argument("--face-mm", action="store_true", help="face 的 X* 单位 mm→m")
    ap.add_argument("--arm-mm", action="store_true", help="arm  的 X* 单位 mm→m")
    ap.add_argument("--out", default=".", help="输出目录")
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)

    face_df = read_motion_ly_or_csv(args.face_file)
    arm_df  = read_motion_ly_or_csv(args.arm_file)

    t_face, face_idx = find_time_column(face_df)
    t_arm,  arm_idx  = find_time_column(arm_df)

    face_df = slice_lines(face_df, args.start_line, args.end_line)
    arm_df  = slice_lines(arm_df,  args.start_line, args.end_line)

# 同步裁剪时间数组
    if args.start_line is not None or args.end_line is not None:
        t_face = t_face[args.start_line:args.end_line]
        t_arm  = t_arm[args.start_line:args.end_line]

    # face_df, t_face = trim(face_df, t_face, args.start, args.end)
    # arm_df,  t_arm  = trim(arm_df,  t_arm,  args.start, args.end)

    # face 期望 X1/Y1/Z1；arm 期望 X3/Y3/Z3，但也允许自动匹配任意编号
    face_xyz = pick_xyz(face_df, ["X1","Y1","Z1"])
    arm_xyz  = pick_xyz(arm_df,  ["X3","Y3","Z3"])

    if args.face_mm: face_xyz = face_xyz / 1000.0
    if args.arm_mm:  arm_xyz  = arm_xyz  / 1000.0
    # 90Hz 下 3Hz → alpha ≈ 0.189，可取 0.18~0.20
    face_xyz = denoise_xyz(face_xyz, do_hampel=True, do_ema=True, alpha=0.19, do_detrend=False)
    arm_xyz  = denoise_xyz(arm_xyz,  do_hampel=True, do_ema=True, alpha=0.19, do_detrend=False)
    # 图1：face
    plt.figure()
    plt.plot(t_face, face_xyz[:,0], label="X_face")
    plt.plot(t_face, face_xyz[:,1], label="Y_face")
    plt.plot(t_face, face_xyz[:,2], label="Z_face")
    plt.xlabel("time (s)" if not face_idx else "sample"); plt.ylabel("value")
    plt.title("face: X/Y/Z"); plt.legend(); plt.tight_layout()
    plt.savefig(os.path.join(args.out, "face_xyz.png")); plt.close()

    # 图2：arm
    plt.figure()
    plt.plot(t_arm, arm_xyz[:,0], label="X_arm")
    plt.plot(t_arm, arm_xyz[:,1], label="Y_arm")
    plt.plot(t_arm, arm_xyz[:,2], label="Z_arm")
    plt.xlabel("time (s)" if not arm_idx else "sample"); plt.ylabel("value")
    plt.title("arm: X/Y/Z"); plt.legend(); plt.tight_layout()
    plt.savefig(os.path.join(args.out, "arm_xyz.png")); plt.close()

    # 图3：3D 叠加
    try:
        from mpl_toolkits.mplot3d import Axes3D  # noqa
        fig = plt.figure(); ax = fig.add_subplot(111, projection="3d")
        ax.plot(face_xyz[:,0], face_xyz[:,1], face_xyz[:,2], label="face")
        ax.plot(arm_xyz[:,0],  arm_xyz[:,1],  arm_xyz[:,2],  label="arm")
        ax.set_xlabel("X"); ax.set_ylabel("Y"); ax.set_zlabel("Z")
        ax.set_title("3D Trajectories"); ax.legend(); plt.tight_layout()
        plt.savefig(os.path.join(args.out, "face_arm_3d.png")); plt.close()
    except Exception as e:
        print("3D 绘图失败（可忽略）：", e)

    # 图4：叠加曲线（同一张图里画 face 与 arm 的 X/Y/Z）
    plt.figure()
    plt.plot(t_face, face_xyz[:, 0], label="face.X")
    plt.plot(t_arm,  arm_xyz[:, 0],  label="arm.X",  linestyle="--")

    plt.plot(t_face, face_xyz[:, 1], label="face.Y")
    plt.plot(t_arm,  arm_xyz[:, 1],  label="arm.Y",  linestyle="--")

    plt.plot(t_face, face_xyz[:, 2], label="face.Z")
    plt.plot(t_arm,  arm_xyz[:, 2],  label="arm.Z",  linestyle="--")

    plt.xlabel("time (s)" if not face_idx else "sample")
    plt.ylabel("value")
    plt.title("face vs arm — X/Y/Z (overlay)")
    plt.legend(ncol=2)
    plt.tight_layout()
    plt.savefig(os.path.join(args.out, "face_arm_overlay.png"))
    plt.close()


    print("完成：face_xyz.png, arm_xyz.png, face_arm_3d.png")

if __name__ == "__main__":
    main()
